package com.sts.sink.impl;

import cn.hutool.core.collection.CollectionUtil;
import com.sts.sink.StsSink;
import com.sts.sink.model.CsvSinkConfig;
import com.sts.source.StsSource;
import com.sts.task.StsTask;
import com.sts.transform.StsTransform;
import de.siegmar.fastcsv.writer.CsvWriter;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.List;
import java.util.stream.Gatherers;
import java.util.stream.Stream;

@Slf4j
public class CsvSinkImpl implements StsSink {

    private final StsSource stsSource;
    private final CsvSinkConfig csvSinkConfig;
    private final StsTask stsTask;

    public CsvSinkImpl(StsSource stsSource, CsvSinkConfig csvSinkConfig, StsTask stsTask) {
        this.stsSource = stsSource;
        this.csvSinkConfig = csvSinkConfig;
        this.stsTask = stsTask;
    }

    @Override
    public void insert() {
        try {
            try (CsvWriter csvWriter = buildCsvWriter()) {
                if (csvSinkConfig.isHasHeader()) {
                    // 导出头
                    List<String> header = stsSource.getHeader();
                    csvWriter.writeRecord(header);
                }

                Stream<List<String>> dataStream = stsSource.getDataStream();

                if (CollectionUtil.isNotEmpty(stsTask.getStsTransformList())) {
                    for (StsTransform<?> stsTransform : stsTask.getStsTransformList()) {
                        dataStream = stsTransform.transform(dataStream);
                    }
                }

                dataStream.takeWhile(_ -> stsTask.isFlag())
                        .gather(Gatherers.windowFixed(csvSinkConfig.getBatchSize()))
                        .forEach(dataList -> {
                            for (List<String> rowValue : dataList) {
                                csvWriter.writeRecord(rowValue);
                            }
                            try {
                                csvWriter.flush();
                            } catch (IOException e) {
                                throw new RuntimeException(e.getMessage(), e);
                            }

                            stsTask.updateTotalSuccessCount(stsTask, dataList.size());
                        });
            } catch (IOException e) {
                throw new RuntimeException(e.getMessage(), e);
            }
        } finally {
            try {
                stsSource.close();
            } catch (Exception e) {
                log.error("关闭连接失败: {}", e.getMessage(), e);
            }
        }
    }

    @Override
    public void upsert() {
        insert();
    }

    private CsvWriter buildCsvWriter() {
        try {
            return CsvWriter.builder()
                    .fieldSeparator(csvSinkConfig.getFieldSeparator())
                    .quoteCharacter(csvSinkConfig.getQuoteCharacter())
                    .lineDelimiter(csvSinkConfig.getLineDelimiter())
                    .bufferSize(csvSinkConfig.getBufferSize())
                    .build(Paths.get(csvSinkConfig.getFilePath()));
        } catch (IOException e) {
            throw new RuntimeException(e.getMessage(), e);
        }
    }
}
